import cv2
import numpy as np
from matplotlib import pyplot as plt


class FlameDetector:
    def __init__(self):
        # 定义火焰颜色的HSV范围（红色和黄色）
        self.lower_flame = np.array([0, 120, 150])
        self.upper_flame = np.array([45, 255, 255])

        # 用于可视化
        self.font = cv2.FONT_HERSHEY_SIMPLEX

    def detect_flame(self, image_path):
        """检测单张图像中的火焰"""
        # 读取图像
        frame = cv2.imread(image_path)
        if frame is None:
            print(f"错误：无法读取图像 {image_path}")
            return None, None, None

        # 存储原始图像用于显示
        original = frame.copy()

        # 转换为HSV颜色空间
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

        # 创建火焰颜色掩码
        flame_mask = cv2.inRange(hsv, self.lower_flame, self.upper_flame)

        # 形态学操作以去除噪声
        kernel = np.ones((5, 5), np.uint8)
        flame_mask = cv2.morphologyEx(flame_mask, cv2.MORPH_OPEN, kernel)
        flame_mask = cv2.morphologyEx(flame_mask, cv2.MORPH_CLOSE, kernel)

        # 寻找轮廓
        contours, _ = cv2.findContours(flame_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # 存储检测到的火焰区域
        flame_regions = []

        # 处理每个轮廓
        for contour in contours:
            # 计算轮廓面积
            area = cv2.contourArea(contour)

            # 忽略太小的区域
            if area < 500:
                continue

            # 计算边界框
            x, y, w, h = cv2.boundingRect(contour)

            # 检查宽高比（火焰通常较高）
            aspect_ratio = h / float(w)
            if aspect_ratio < 0.8:
                continue

            # 检查区域亮度变化（火焰中心通常更亮）
            region = frame[y:y + h, x:x + w]
            gray_region = cv2.cvtColor(region, cv2.COLOR_BGR2GRAY)
            _, thresh = cv2.threshold(gray_region, 200, 255, cv2.THRESH_BINARY)
            bright_area = np.count_nonzero(thresh)

            if bright_area / area < 0.1:
                continue

            # 添加到火焰区域列表
            flame_regions.append((x, y, w, h))

            # 在原始图像上绘制边界框和轮廓
            cv2.rectangle(original, (x, y), (x + w, y + h), (0, 0, 255), 2)
            cv2.drawContours(original, [contour], -1, (0, 255, 255), 2)
            cv2.putText(original, "FLAME", (x, y - 10),
                        self.font, 0.7, (0, 0, 255), 2)

        # 在图像上显示状态
        status = "Flame Detected!" if flame_regions else "No Flame"
        color = (0, 0, 255) if flame_regions else (0, 255, 0)
        cv2.putText(original, f"Status: {status}", (10, 30),
                    self.font, 1, color, 2)

        # 创建彩色掩码用于可视化
        color_mask = cv2.cvtColor(flame_mask, cv2.COLOR_GRAY2BGR)
        color_mask[flame_mask == 255] = (0, 0, 255)  # 将掩码中的白色区域转为红色

        return original, color_mask, flame_regions


def display_results(original, mask, results):
    """显示检测结果"""
    plt.figure(figsize=(15, 10))

    # 显示原始图像
    plt.subplot(1, 2, 1)
    plt.imshow(cv2.cvtColor(original, cv2.COLOR_BGR2RGB))
    plt.title('Flame Detection Result')
    plt.axis('off')

    # 显示火焰掩码
    plt.subplot(1, 2, 2)
    plt.imshow(cv2.cvtColor(mask, cv2.COLOR_BGR2RGB))
    plt.title('Flame Color Mask')
    plt.axis('off')

    plt.tight_layout()
    plt.show()

    # 打印检测结果
    if results:
        print(f"检测到 {len(results)} 个火焰区域:")
        for i, (x, y, w, h) in enumerate(results, 1):
            print(f"区域 {i}: 位置({x}, {y}), 宽度={w}, 高度={h}")
    else:
        print("未检测到火焰")


def main():
    # 初始化火焰检测器
    detector = FlameDetector()

    # 图像文件路径
    image_path = "../files/fire_stairs.webp"  # 替换为你的图像路径

    # 检测火焰
    result_img, mask_img, flame_regions = detector.detect_flame(image_path)

    if result_img is not None:
        # 显示结果
        display_results(result_img, mask_img, flame_regions)

        # 保存结果
        cv2.imwrite("detection_result.jpg", result_img)
        cv2.imwrite("flame_mask.jpg", mask_img)
        print("结果已保存为 'detection_result.jpg' 和 'flame_mask.jpg'")


if __name__ == "__main__":
    main()